#!/usr/bin/env python3
"""
周期性提示生成器

基于心脏周期性的提示生成模块，利用超声心动图的时序特性
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Tuple, Optional, Any

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from .base_prompt import BasePromptGenerator
from utils.logger import get_logger


class PeriodicPromptGenerator(BasePromptGenerator):
    """周期性提示生成器"""
    
    def __init__(self, 
                 input_dim: int = 256,
                 prompt_dim: int = 256,
                 sequence_length: int = 20,
                 num_periods: int = 4,
                 device: str = "cuda"):
        """
        初始化周期性提示生成器
        
        Args:
            input_dim: 输入特征维度
            prompt_dim: 提示特征维度
            sequence_length: 序列长度
            num_periods: 周期数量
            device: 设备
        """
        super().__init__(input_dim, prompt_dim, device)
        
        self.sequence_length = sequence_length
        self.num_periods = num_periods
        
        # 时序编码器
        self.temporal_encoder = nn.LSTM(
            input_size=input_dim,
            hidden_size=prompt_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.1
        )
        
        # 周期检测器
        self.period_detector = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, num_periods),
            nn.Softmax(dim=-1)
        )
        
        # 相位编码器
        self.phase_encoder = nn.Sequential(
            nn.Linear(prompt_dim + num_periods, prompt_dim),
            nn.ReLU(),
            nn.Linear(prompt_dim, prompt_dim)
        )
        
        # 提示生成器
        self.point_generator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 2)  # (x, y)
        )
        
        self.box_generator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 4)  # (x1, y1, x2, y2)
        )
        
        self.mask_generator = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 32 * 32)  # 32x32掩码
        )
        
        # 注意力机制
        self.attention = nn.MultiheadAttention(
            embed_dim=prompt_dim,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        )
        
        self.logger.info(f"周期性提示生成器初始化完成: {self.get_model_info()}")
    
    def detect_cardiac_cycle(self, 
                           temporal_features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        检测心脏周期
        
        Args:
            temporal_features: 时序特征 (B, T, C)
            
        Returns:
            (周期概率, 相位信息)
        """
        # 检测周期
        period_probs = self.period_detector(temporal_features)  # (B, T, num_periods)
        
        # 计算相位
        phase_info = self.phase_encoder(
            torch.cat([temporal_features, period_probs], dim=-1)
        )  # (B, T, prompt_dim)
        
        return period_probs, phase_info
    
    def generate_periodic_prompts(self, 
                                phase_info: torch.Tensor,
                                period_probs: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        生成周期性提示
        
        Args:
            phase_info: 相位信息 (B, T, C)
            period_probs: 周期概率 (B, T, num_periods)
            
        Returns:
            提示字典
        """
        batch_size, seq_len, _ = phase_info.shape
        
        # 生成点提示
        points = self.point_generator(phase_info)  # (B, T, 2)
        
        # 生成框提示
        boxes = self.box_generator(phase_info)  # (B, T, 4)
        
        # 生成掩码提示
        mask_logits = self.mask_generator(phase_info)  # (B, T, 32*32)
        masks = mask_logits.view(batch_size, seq_len, 32, 32)
        
        # 根据周期概率调整提示强度
        period_weights = period_probs.max(dim=-1)[0].unsqueeze(-1)  # (B, T, 1)
        
        points = points * period_weights
        boxes = boxes * period_weights
        masks = masks * period_weights.unsqueeze(-1).unsqueeze(-1)
        
        return {
            'points': points,
            'boxes': boxes,
            'masks': masks,
            'period_probs': period_probs,
            'phase_info': phase_info
        }
    
    def forward(self, 
                image_features: torch.Tensor,
                context: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            context: 上下文信息，包含时序信息
            
        Returns:
            提示字典
        """
        batch_size = image_features.shape[0]
        
        if context and 'temporal_sequence' in context:
            # 使用提供的时序序列
            temporal_sequence = context['temporal_sequence']  # (B, T, C)
        else:
            # 从当前图像特征生成时序序列
            # 这里简化处理，实际应该使用历史帧
            temporal_sequence = image_features.unsqueeze(1).repeat(1, self.sequence_length, 1, 1, 1)
            temporal_sequence = temporal_sequence.view(batch_size, self.sequence_length, -1)
        
        # 时序编码
        temporal_output, _ = self.temporal_encoder(temporal_sequence)  # (B, T, prompt_dim)
        
        # 检测心脏周期
        period_probs, phase_info = self.detect_cardiac_cycle(temporal_output)
        
        # 生成周期性提示
        prompts = self.generate_periodic_prompts(phase_info, period_probs)
        
        # 添加注意力机制
        attended_features, _ = self.attention(
            temporal_output, temporal_output, temporal_output
        )
        
        # 融合注意力特征
        prompts['attended_features'] = attended_features
        
        return prompts
    
    def generate_prompts(self, 
                        image: torch.Tensor,
                        metadata: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        生成提示
        
        Args:
            image: 输入图像 (B, C, H, W)
            metadata: 元数据，包含时序信息
            
        Returns:
            提示字典
        """
        # 提取图像特征
        image_features = self.extract_features(image)
        
        # 构建上下文
        context = {}
        if metadata:
            context.update(metadata)
        
        # 生成提示
        prompts = self.forward(image_features, context)
        
        return prompts
    
    def compute_periodic_loss(self, 
                            predictions: Dict[str, torch.Tensor],
                            targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        计算周期性损失
        
        Args:
            predictions: 预测结果
            targets: 目标结果
            
        Returns:
            损失值
        """
        total_loss = 0.0
        
        # 基础提示损失
        total_loss += self.compute_loss(predictions, targets)
        
        # 周期性一致性损失
        if 'period_probs' in predictions and 'period_probs' in targets:
            period_loss = F.kl_div(
                F.log_softmax(predictions['period_probs'], dim=-1),
                F.softmax(targets['period_probs'], dim=-1),
                reduction='batchmean'
            )
            total_loss += 0.1 * period_loss
        
        # 相位平滑损失
        if 'phase_info' in predictions:
            phase_info = predictions['phase_info']  # (B, T, C)
            phase_diff = phase_info[:, 1:] - phase_info[:, :-1]
            smooth_loss = torch.mean(torch.norm(phase_diff, dim=-1))
            total_loss += 0.01 * smooth_loss
        
        return total_loss
    
    def get_cycle_analysis(self, 
                          temporal_sequence: torch.Tensor) -> Dict[str, Any]:
        """
        获取周期分析结果
        
        Args:
            temporal_sequence: 时序序列 (B, T, C)
            
        Returns:
            周期分析结果
        """
        with torch.no_grad():
            # 时序编码
            temporal_output, _ = self.temporal_encoder(temporal_sequence)
            
            # 检测周期
            period_probs, phase_info = self.detect_cardiac_cycle(temporal_output)
            
            # 分析结果
            dominant_periods = torch.argmax(period_probs, dim=-1)  # (B, T)
            cycle_confidence = torch.max(period_probs, dim=-1)[0]  # (B, T)
            
            return {
                'dominant_periods': dominant_periods.cpu().numpy(),
                'cycle_confidence': cycle_confidence.cpu().numpy(),
                'period_probs': period_probs.cpu().numpy(),
                'phase_info': phase_info.cpu().numpy()
            }
